#!/usr/bin/python

# Copyright (C) 2009 ApplianSys Ltd.
# 
# This program is free software; you can redistribute it and/or
# modify it under the terms of the GNU General Public License
# as published by the Free Software Foundation; either version 2
# of the License, or (at your option) any later version.
# 
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
# 
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, 
# MA  02110-1301, USA.

"""
This file contains the definition of the worthyd udp broadcast packet
structure.
"""

import struct
import md5
import sys

import worthyd

FORMAT_FH = ("!6sBB"         # protocol_id, version and ip_version
             "xxxxxxxxxxxxL" # worthiness (with room to be extended)
             "L"             # num_members
             "16s"           # members_md5
             "HH"            # valid_for, len_group_name
             )

FORMAT_SH = "!16s"    # md5 checksum

FIXED_LEN = struct.calcsize(FORMAT_FH)
PROTOCOL_ID = "worthy"
PROTOCOL_VERSION = 0
LEN_MD5_SUM = 16
         

class PacketException(Exception):
    pass


class BroadcastPacket(object):
    """
    This class represents a broadcast packet, either send or recieved.
    It handles all the encoding/decodeing and packet validation.

    Test to check that we can encode and decode a packet successfully.

    >>> packet = BroadcastPacket(worthiness=7, 
    ...                          num_members=3,
    ...                          members_md5='asdfghjkqwertyui',
    ...                          valid_for=20,
    ...                          group_name='group_foobar',
    ...                          )
    >>> packet.encode()
    >>> packet.buffer
    'worthy\\x00\\x04\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x07\\x00\\x00\\x00\\x03asdfghjkqwertyui\\x00\\x14\\x00\\x0cgroup_foobar\\x0c`\\xbb\\xccr\\xe8\\xac\\xc0\\xbf\\xe9\\xcb\\xa3\\xbc\\xca\\x8b\\x97'
    >>> buffer = packet.buffer
    >>> recvd_packet = BroadcastPacket(buffer=buffer)
    >>> recvd_packet.worthiness
    7
    >>> recvd_packet.num_members
    3
    >>> recvd_packet.members_md5
    'asdfghjkqwertyui'
    >>> recvd_packet.valid_for
    20
    >>> recvd_packet.group_name
    'group_foobar'


    Test that damaged packet is rejected.

    >>> damaged_buffer = buffer[:-1] + 'z'
    >>> recvd_packet = BroadcastPacket(buffer=damaged_buffer)
    Traceback (most recent call last):
        ...
    PacketException: Incorrect checksum. '\\x0c`\\xbb\\xccr\\xe8\\xac\\xc0\\xbf\\xe9\\xcb\\xa3\\xbc\\xca\\x8b\\x97' != '\\x0c`\\xbb\\xccr\\xe8\\xac\\xc0\\xbf\\xe9\\xcb\\xa3\\xbc\\xca\\x8bz'


    Check that it can cope with extremely short input.

    >>> damaged_buffer = buffer[:10]
    >>> recvd_packet = BroadcastPacket(buffer=damaged_buffer)
    Traceback (most recent call last):
        ...
    PacketException: Extremely short packet detected.


    Check that it can cope with short input.

    >>> damaged_buffer = buffer[:-1]
    >>> recvd_packet = BroadcastPacket(buffer=damaged_buffer)
    Traceback (most recent call last):
        ...
    PacketException: Bad packet length detected.


    Check that it can cope with appended junk.

    >>> damaged_buffer = buffer[:] + "hellow world"
    >>> recvd_packet = BroadcastPacket(buffer=damaged_buffer)
    Traceback (most recent call last):
        ...
    PacketException: Bad packet length detected.


    """

    def __init__(self,
                 protocol_id=PROTOCOL_ID,
                 protocol_version=PROTOCOL_VERSION,
                 ip_version=4,
                 worthiness=0,
                 num_members=0,
                 members_md5='',
                 valid_for=0,
                 group_name='',
                 buffer=None,
                 sender=None,
                 ):
        """
        @param protocol_id: always 'worthyd'
        @param protocol_version: always 0
        @param ip_version: either 4 or 6 (6 not implemented yet)
        @param worthiness: the worthiness of the sending node
        @type worthiness: Integer
        @param num_members: the number of group memebrs about which this node
                            knows
        @param members_md5: the md5sum of the addresses of the members of the 
                            group, in ascending big-endian order 
        @param valid_for: how long receivers of this packet should believe 
                          that this node is alive, if no later packet is 
                          received from this node
        @param group_name: the name of the group
        @buffer: the representation of the udp data
        @sender: 
        """
        self.protocol_id = protocol_id
        self.protocol_version = protocol_version
        self.ip_version = ip_version
        self.worthiness = worthiness
        self.num_members = num_members
        self.members_md5 = members_md5
        self.valid_for = valid_for
        self.group_name = group_name
        self.len_group_name = len(group_name)
        self.buffer = buffer
        self.sender = sender
        
        # if we've been given a buffer, decode it
        if buffer != None:
            self.decode()
        else:
            self.encode()


    def encode_fh(self):
        """
        This encodes the first_half fields fo the packet into self.buffer.

        @return: all of the fields, except the md5checksum, encoded into 
                 network order bytes.
        """

        first_half = struct.pack(FORMAT_FH, 
                          self.protocol_id,
                          self.protocol_version,
                          self.ip_version,
                          self.worthiness,
                          self.num_members,
                          self.members_md5,
                          self.valid_for,
                          self.len_group_name,
                          ) + self.group_name

        return first_half


    def encode(self):
        """
        Encodes the fields into self.buffer.

        All of the fields, including the md5checksum, encoded into 
        network order bytes.
        """

        first_half = self.encode_fh()
        checksum = md5.new(first_half).digest()
        self.buffer = self.encode_fh() + struct.pack(FORMAT_SH, checksum)


    def decode(self):
        """
        This decodes self.buffer into the fields of the packet.
        """
        try:
            (self.protocol_id,
             self.protocol_version,
             self.ip_version,
             self.worthiness,
             self.num_members,
             self.members_md5,
             self.valid_for,
             self.len_group_name) = struct.unpack(FORMAT_FH, 
                                                  self.buffer[:FIXED_LEN],
                                                  )
        except Exception, e:
            raise PacketException("Extremely short packet detected.") 

        if self.protocol_id != PROTOCOL_ID:
            raise PacketException("Incorrect protocol_id of %s." 
                                  % (self.protocol_id)) 
        if self.protocol_version != PROTOCOL_VERSION:
            raise PacketException("Incorrect protocol_version of %s." 
                                  % (self.protocol_version)) 
            
        if self.len_group_name > worthyd.MAX_NAME_LEN:
            raise PacketException("Incorrect max_name_len exceeded %s."
                                  % (self.len_group_name))        

        expected_data_length = FIXED_LEN + self.len_group_name
        expected_full_length = expected_data_length + LEN_MD5_SUM

        if len(self.buffer) != expected_full_length:
            raise PacketException("Bad packet length detected.") 

        self.group_name = self.buffer[FIXED_LEN:expected_data_length]

        try:
            self.checksum = struct.unpack(FORMAT_SH, 
                                          self.buffer[expected_data_length 
                                                      : 
                                                      expected_full_length
                                                      ]
                                          )[0]
        except Exception, e:
            # WARN: I think this should never be reached, as it will have
            # already thrown a 'bad packet length' before here.
            # Leaving it in for belt and braces.
            raise PacketException("Short packet detected.") 

        expected_checksum = \
                md5.new(self.buffer[:expected_data_length]).digest()

        if expected_checksum != self.checksum:
            raise PacketException("Incorrect checksum. %s != %s" % 
                (expected_checksum.__repr__(), 
                 self.checksum.__repr__(), 
                 ))


def _test():
    import doctest
    failures, x = doctest.testmod()
    if failures > 0:
        sys.exit(1)


if __name__ == "__main__":
    _test()

 
